# server/src/knowledge/postgres_provider.py
# PostgreSQL实现
import os
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from .models import Document, KnowledgeBaseItem, Base
from .file_parser import FileParser
from .exceptions import FileParseError

class PostgresKnowledgeProvider(KnowledgeProvider):
    def __init__(self, db_url: str):
        self.engine = create_async_engine(db_url)
        self.parser = FileParser()

    async def initialize(self):
        """初始化数据库"""
        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

    async def add_document(self, file_content: bytes, filename: str, metadata: dict) -> int:
        """添加文档到知识库"""
        try:
            # 解析文件内容
            content = self.parser.parse_file(file_content, filename)
            file_type = os.path.splitext(filename)[-1].lower().strip('.')
            
            async with AsyncSession(self.engine) as session:
                doc = Document(
                    filename=filename,
                    file_type=file_type,
                    content=content,
                    metadata=metadata
                )
                session.add(doc)
                await session.commit()
                return doc.id
        except Exception as e:
            raise FileParseError(f"文件解析失败: {str(e)}")

    async def search(self, query: str, top_k: int=3) -> List[str]:
        """简单文本搜索（可扩展为向量搜索）"""
        async with AsyncSession(self.engine) as session:
            stmt = select(Document.content).where(
                Document.is_deleted == 0
            ).order_by(Document.updated_at.desc()).limit(top_k)
            
            result = await session.execute(stmt)
            return [row[0] for row in result.scalars().all()]

    async def delete_document(self, doc_id: int) -> bool:
        """软删除文档"""
        async with AsyncSession(self.engine) as session:
            stmt = update(Document).where(
                Document.id == doc_id
            ).values(is_deleted=1)
            
            await session.execute(stmt)
            await session.commit()
            return True

    async def list_documents(self, page: int=1, page_size: int=10) -> List[dict]:
        """列出所有文档"""
        async with AsyncSession(self.engine) as session:
            stmt = select(Document).where(
                Document.is_deleted == 0
            ).offset((page-1)*page_size).limit(page_size)
            
            result = await session.execute(stmt)
            return [{
                "id": doc.id,
                "filename": doc.filename,
                "file_type": doc.file_type,
                "created_at": doc.created_at.isoformat()
            } for doc in result.scalars().all()]